package main

import (
    "regexp"
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"flag"
	"fmt"
	"time"
	"log"
	"strings"
	"net/http"
	"net/url"
    "crypto/tls"

	"github.com/golang/glog"
	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
    "github.com/samuel/go-zookeeper/zk"

	"http-gateway/powermarket"
	"http-gateway/powermaster"
)

var path_array = [...]string{"/v2/powermaster/user/login", "/v2/powermaster/user","/v2/powermaster/sysuser/login","/v2/powermaster/sysuser","/v2/powermaster/safe","/v2/powermaster/organ","/v2/powermaster","/v2/powermarket","/v2/predict/","/v2/predicting"}
var secret_key = "heiyu secret"

func generateKey() []byte {
	strbytes := []byte(secret_key)
	encodedKey := base64.StdEncoding.EncodeToString(strbytes)
	iv := "0123456789heiyud"
	result, err := AESEncrypt(encodedKey, iv, iv)
	if err != nil {
		log.Println(err)
	}
	return result
}

func AESEncrypt(content, password, iv string) (result []byte, err error) {
	var block cipher.Block
	if block, err = aes.NewCipher([]byte(password)); err != nil {
		log.Println(err)
		return
	}
	encrypt := cipher.NewCBCEncrypter(block, []byte(iv))
	var source []byte = PKCS5Padding([]byte(content), 16)
	var dst []byte = make([]byte, len(source))
	encrypt.CryptBlocks(dst, source)
	result = dst
	return
}

func PKCS5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func run(e *echo.Echo, masterAddr, marketAddr string, port int) error {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	grpc_mux := runtime.NewServeMux(runtime.WithMarshalerOption(runtime.MIMEWildcard, &runtime.JSONPb{OrigName: true, EmitDefaults: true}))
	opts := []grpc.DialOption{grpc.WithInsecure(), grpc.WithMaxMsgSize(20* 1024 * 1024 * 1024)}
	err := powermaster.RegisterPowerMasterServiceHandlerFromEndpoint(ctx, grpc_mux, masterAddr, opts)
	if err != nil {
		return err
	}
	err = powermarket.RegisterPowerMarketServiceHandlerFromEndpoint(ctx, grpc_mux, marketAddr, opts)
	if err != nil {
		return err
	}

	//AI group
	url1, err := url.Parse("http://127.0.0.1:25000")
	if err != nil {
		e.Logger.Fatal(err)
	}
	targets := []*middleware.ProxyTarget{
		{
			URL: url1,
		},
	}
	e.Group("/v2/models/inference/", middleware.Proxy(middleware.NewRoundRobinBalancer(targets)))

	url11, err := url.Parse("http://127.0.0.1:25000")
	if err != nil {
		e.Logger.Fatal(err)
	}
	targets11 := []*middleware.ProxyTarget{
		{
			URL: url11,
		},
	}
	e.Group("/v2/predict/*", middleware.Proxy(middleware.NewRoundRobinBalancer(targets11)))

	url12, err := url.Parse("http://127.0.0.1:25000")
	if err != nil {
		e.Logger.Fatal(err)
	}
	targets12 := []*middleware.ProxyTarget{
		{
			URL: url12,
		},
	}
	e.Group("/v2/predicting/*", middleware.Proxy(middleware.NewRoundRobinBalancer(targets12)))

	//AC_Control
	url3, err := url.Parse("http://127.0.0.1:28000")
	if err != nil {
		e.Logger.Fatal(err)
	}
	targets2 := []*middleware.ProxyTarget{
		{
			URL: url3,
		},
	}
	e.Group("/v2/models/training/*", middleware.Proxy(middleware.NewRoundRobinBalancer(targets2)))

    //I6000 group
    url2, err := url.Parse("http://127.0.0.1:29000")
    if err != nil {
        e.Logger.Fatal(err)
    }
    targets1 := []*middleware.ProxyTarget{
       	{
        	URL: url2,
       	},
    }
    e.Group("/v2/i6000/*", middleware.Proxy(middleware.NewRoundRobinBalancer(targets1)))

    //iscLogin group
        url5, err := url.Parse("http://127.0.0.1:28010")
        if err != nil {
            e.Logger.Fatal(err)
        }
        targets5 := []*middleware.ProxyTarget{
           	{
            	URL: url5,
           	},
        }
        e.Group("/v2/iscLogin/*", middleware.Proxy(middleware.NewRoundRobinBalancer(targets5)))

	//backend group
	//jwt middleware for authorization
	backend := e.Group("/v2/*")
	backend.Use(middleware.JWTWithConfig(middleware.JWTConfig{ TokenLookup: "query:token",
		SigningKey:  generateKey(),
		Skipper: func(c echo.Context) bool { //skip login
			path := c.Request().URL.Path
			for index := range path_array {
				if strings.Contains(path, path_array[index]) {
					return true
				}
			}
			if strings.Contains(c.Request().Host, "localhost") { //skip Host:localhost/127.0.0.1
				return true
			}
			if strings.Contains(c.Request().Host, "127.0.0.1") {
				return true
			}
			return true
		},
	}))
	backend.Any("", echo.WrapHandler(grpc_mux))

	//Static
	e.Static("/", "fe")

	e.Start(fmt.Sprintf(":%d", port))
	return nil
}

func redirectHTTPS(port int) echo.MiddlewareFunc {
    return func(next echo.HandlerFunc) echo.HandlerFunc {
        return func(c echo.Context) error {
            req, scheme := c.Request(), c.Scheme()
            host := strings.Split(req.Host, ":")[0]
            url := fmt.Sprintf("https://%s:%d%s", host, port, req.RequestURI)
            if scheme != "https" {
                return c.Redirect(http.StatusMovedPermanently, url)
            }

            return next(c)
        }
    }
}

func startTLS(e *echo.Echo, address string, certFile, keyFile string) (err error) {

	if certFile == "" || keyFile == "" {
	    return
	}

	s := e.TLSServer
	s.TLSConfig = new(tls.Config)
	s.TLSConfig.Certificates = make([]tls.Certificate, 1)
	if s.TLSConfig.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile); err != nil {
		return
	}
	s.TLSConfig.MinVersion = tls.VersionTLS12

	//return e.startTLS(address)
	return startTLS1(e,address)
}

func startTLS1(e *echo.Echo, address string) error {
	s := e.TLSServer
	s.Addr = address
	if !e.DisableHTTP2 {
		s.TLSConfig.NextProtos = append(s.TLSConfig.NextProtos, "h2")
	}
	return e.StartServer(e.TLSServer)
}

func runTLS(e *echo.Echo, port int, perm, key string) {
    e.Use(redirectHTTPS(port))
    startTLS(e,fmt.Sprintf(":%d", port), perm, key)
}

func xffClean(next echo.HandlerFunc) echo.HandlerFunc {
    return func(c echo.Context) error {
        c.Request().Header.Del(echo.HeaderXForwardedFor)
        return next(c)
    }
}

func judgeIp() echo.MiddlewareFunc {
   return func(next echo.HandlerFunc) echo.HandlerFunc {
       return func(c echo.Context) error {
        req := c.Request()
			var ip string = req.RemoteAddr
			var result int =  strings.Index(ip,"30")
			log.Printf("getway ip地址:%s",ip);
	        if result != 0 {
				log.Printf("getway ip不在内网内:%s",ip);
                return c.JSON(http.StatusForbidden,"Not in Intranet Environment")
           }
           return next(c)
        }
    }
}

func filter() echo.MiddlewareFunc {
    return func(next echo.HandlerFunc) echo.HandlerFunc {
        return func(c echo.Context) error {
            formParams, _ := c.FormParams()
            for key,value:=range formParams{
                formParams[key] = SanitizeName(value)
            }
            c.Request().Form = formParams

            queryString := c.QueryString()
            reNameBlacklist.ReplaceAllString(queryString, "")
            c.Request().URL.RawQuery = queryString

            return next(c)
        }
    }
}

var reNameBlacklist = regexp.MustCompile(`(~|!|@|#|$|%|\^|\$|\.|\*|\+|\?|\{|\}|\(|\)|\[|\]|,|=|-|_|《|》|>|<|\/|:|\n|\r|\a|\f|\t|\v|\b|。|，|、|‘|；|：|“|”|【|】|-|？｜￥|（|）|'|"|\|')*`)

var sqlBlockList= regexp.MustCompile(`(?:')|(?:--)|(/\\*(?:.|[\\n\\r])*?\\*/)|(\b(select|update|and|or|delete|insert|trancate|char|chr|into|substr|ascii|declare|exec|count|master|into|drop|execute)\b)`)

// replace ""
func SanitizeName(names []string) []string {
    for index,_ := range names {
        //replace ""
        names[index] = strings.Replace(names[index], "alert", "", -1)
        names[index] = strings.Replace(names[index], "anchor", "", -1)
        names[index] = strings.Replace(names[index], "assign", "", -1)
        names[index] = strings.Replace(names[index], "anchor", "", -1)
        names[index] = strings.Replace(names[index], "blur", "", -1)
        names[index] = strings.Replace(names[index], "button", "", -1)
        names[index] = strings.Replace(names[index], "close", "", -1)
        names[index] = strings.Replace(names[index], "confirm", "", -1)
        names[index] = strings.Replace(names[index], "crypto", "", -1)
        names[index] = strings.Replace(names[index], "document", "", -1)
        names[index] = strings.Replace(names[index], "element", "", -1)
        names[index] = strings.Replace(names[index], "embed", "", -1)
        names[index] = strings.Replace(names[index], "event", "", -1)
        names[index] = strings.Replace(names[index], "form", "", -1)
        names[index] = strings.Replace(names[index], "navigate", "", -1)
        names[index] = strings.Replace(names[index], "option", "", -1)
        names[index] = strings.Replace(names[index], "layer", "", -1)
        names[index] = strings.Replace(names[index], "reset", "", -1)
        names[index] = strings.Replace(names[index], "prompt", "", -1)
        names[index] = strings.Replace(names[index], "plugin", "", -1)
        names[index] = strings.Replace(names[index], "text", "", -1)
        names[index] = strings.Replace(names[index], "link", "", -1)
        names[index] = strings.Replace(names[index], "packages", "", -1)
        names[index] = strings.Replace(names[index], "window", "", -1)
        names[index] = strings.Replace(names[index], "constructor", "", -1)
        names[index] = strings.Replace(names[index], "decodeURI", "", -1)
        names[index] = strings.Replace(names[index], "location", "", -1)
        names[index] = strings.Replace(names[index], "[s|S][c|C][r|R][i|C][p|P][t|T]", "", -1)
        names[index] = strings.Replace(names[index], "[e|E][v|V][a|A][l|L]\\((.*)\\)", "", -1)
        names[index] = strings.Replace(names[index], "[\\s]*[j|J][a|A][v|V][a|A][s|S][c|C][r|R][i|I][p|P][t|T]:(.*)", "", -1)
        names[index] = strings.Replace(names[index], "`", "", -1)
        names[index] = strings.Replace(names[index], "&", "", -1)
        names[index] = strings.Replace(names[index], "·", "", -1)
        names[index] = strings.Replace(names[index], "～", "", -1)
        names[index] = strings.Replace(names[index], "！", "", -1)
        names[index] = strings.Replace(names[index], "@", "", -1)
        names[index] = strings.Replace(names[index], "#", "", -1)
        names[index] = strings.Replace(names[index], "¥", "", -1)
        names[index] = strings.Replace(names[index], "%", "", -1)
        names[index] = strings.Replace(names[index], "……", "", -1)
        names[index] = strings.Replace(names[index], "（", "", -1)
        names[index] = strings.Replace(names[index], "）", "", -1)
        names[index] = strings.Replace(names[index], "「", "", -1)
        names[index] = strings.Replace(names[index], "」", "", -1)
        names[index] = strings.Replace(names[index], "｜", "", -1)
        names[index] = strings.Replace(names[index], "、", "", -1)
        names[index] = strings.Replace(names[index], "【", "", -1)
        names[index] = strings.Replace(names[index], "】", "", -1)
        names[index] = strings.Replace(names[index], "；", "", -1)
        names[index] = strings.Replace(names[index], "：", "", -1)
        names[index] = strings.Replace(names[index], "。", "", -1)
        names[index] = strings.Replace(names[index], "/", "", -1)
        names[index] = strings.Replace(names[index], "？", "", -1)
        names[index] = strings.Replace(names[index], "，", "", -1)
        names[index] = strings.Replace(names[index], ";", "", -1)
        names[index] = strings.Replace(names[index], "|", "", -1)
        names[index] = strings.Replace(names[index], "——", "", -1)


        reNameBlacklist.ReplaceAllString(names[index], "")
        sqlBlockList.ReplaceAllString(names[index], "")
    }
    return names
}

type tlsConfig struct {
    port int
    permFile string
    keyFile string
}

func zkListen(c chan string, hosts []string) {
    current := hosts[0]

    for {
        stats, _ := zk.FLWSrvr(hosts, 5 * time.Second)

        if stats == nil || len(stats) == 0 {
            glog.Error("fail to get server status")
            time.Sleep(5 * time.Second)
            continue
        }

        for i, stat := range(stats) {
            if stat.Mode == zk.ModeLeader || stat.Mode == zk.ModeStandalone {
                host := fmt.Sprintf("%s:30031",strings.Split(hosts[i], ":")[0])
                if current != host {
                    c <- host
                    current = host
                }

                break
            }
        }

        time.Sleep(5 * time.Second)
    }
}

func serve(sec bool, master, market string, port int, tlsconfig *tlsConfig) *echo.Echo {
	e := echo.New()
	e.Pre(middleware.Rewrite(map[string]string{
      "/v1/*":             "/v2/$1",
      "/v1/predict/*":     "/v2/predict/*",
    }))
	e.HideBanner = true

    if (sec){
        //e.Logger.Fatal("enabling sec...")
         glog.Info("enabling sec...")
         e.Use(Secure())
         e.Use(filter())
     }

	//root middleware
	e.Use(middleware.Recover())
	e.Use(middleware.Logger())
    e.Use(xffClean)

    if nil != tlsconfig {
        glog.Info("enabling SSL...")
        go runTLS(e, tlsconfig.port, tlsconfig.permFile, tlsconfig.keyFile)
    }

	go run(e, master, market, port)

    return e
}

func main() {
	defer glog.Flush()

	masterAddr := flag.String("powermaster", "localhost:30031", "power master service")
	marketAddr := flag.String("powermarket", "localhost:29527", "power market service")
	port := flag.Int("port", 80, "http service port")
    haport := flag.Int("haport", 22181, "HA port of powermasters")

    ssl := flag.Bool("ssl", true, "enable ssl")
    sslPort := flag.Int("sslport", 443, "ssl port")
    sslPerm := flag.String("sslperm", "../conf/alipower.pem", "certificate file")
    sslKey := flag.String("sslkey", "../conf/alipower.key", "certificate key file")

    sec := flag.Bool("sec", false, "disable xss、sql")

	flag.Parse()

    var tlsconfig *tlsConfig
    if *ssl {
        tlsconfig = &tlsConfig{*sslPort, *sslPerm, *sslKey}
    }

    c := make(chan string)

    masters := strings.Split(*masterAddr, ",")
    markets := strings.Split(*marketAddr, ",")

    if len(masters) < 3 {
        // standalone mode
        go func() {
            c <- masters[0]
        }()
    } else {
        // HA mode
        var hosts []string
        for _, master := range(masters) {
            hosts = append(hosts, fmt.Sprintf("%s:%d", strings.Split(master, ":")[0], *haport))
        }
        go zkListen(c, hosts)
    }

    var e *echo.Echo
    for {

        master := <- c
        glog.Info("use backend powermaster ", master)

        if nil != e {
            // gracefully shutdown. do not use e.Close()
            ctx, cancel := context.WithTimeout(context.Background(), 2 * time.Second)
            defer cancel()
            if err := e.Shutdown(ctx); err != nil {
                e.Logger.Fatal(err)
            }
        }

        e = serve(*sec, master, markets[0], *port, tlsconfig)
    }
}
